
from numpy.linalg import norm
from numpy import sqrt, exp, pi, ones, zeros
from numpy import asarray, c_, isscalar

from utils.FemBases import LagrangeP1, LagrangeP2
from utils.Quadrature import LegendreGaussQuadrature
from utils.FemToReal import FemToReal

from models.PIDEStatic1D import PIDEStatic1D
from models.PIDESolver import PIDESolver

from solvers.IMEXTimeStepper import IMEXTimeStepper

import matplotlib.pyplot as plt

import unittest


class PIDETestSuite(unittest.TestCase):

    """ base unit test class. """

    def Kuo_kernel(self, points):
        """ kuo's model, kernel is a simple gaussian. """

        return exp((points - self.m_mu) ** 2 / (2. * self.m_sigma)) \
            / (sqrt(2. * pi) * self.m_sigma)


    """
        kernel 1
            u = 1 | conv = xr - xl
            u = x | conv = ((xr - x) **2 - (xl - x) ** 2) / 2.
            u = x^2 | conv = ((xr - x) ** 3 - (xl - x) ** 3) / 3.

        kernel x
            u = 1 | conv =  x / 2.
            u = x | conv =  -2/3
            u = x^2 | conv =  2*y/3

        kernel x** 2
            u = 1 | conv =  (xr ** 3 - xl ** 3) / 3.
            u = x | conv =  -4*y/3
            u = x^2 | conv =  2*y**2/3 + 2/5
    """

    def exact_kernel(self, points):
        """ suite of kernels. """

        if self.m_kernel == '1':
            return ones(points.shape)
        elif self.m_kernel == 'x':
            return points
        elif self.m_kernel == 'x2':
            return points ** 2
        assert False

    @classmethod
    def x0_exact(cls, points):
        """ suite of exact solutions. """

        return ones(points.shape)

    def x0_load(self, points):
        """ PDE x0_load. """

        pde = self.m_reac - self.m_lamb
        if self.m_kernel == '1':
            pde += ones(points.shape) * self.m_lamb * 2
        if self.m_kernel == 'x':
            pde += self.m_lamb * 2 * points
        elif self.m_kernel == 'x2':
            pde += self.m_lamb * 2 * (points ** 2 + ones(points.shape)
                    / 3.)
        return pde

    @classmethod
    def x1_exact(cls, points):
        """ suite of exact solutions. """

        return points

    def x1_load(self, points):
        """ PDE x0_load. """

        pde = self.m_conv + (self.m_reac - self.m_lamb) * points
        if self.m_kernel == '1':
            return pde
        elif self.m_kernel == 'x':
            return pde - self.m_lamb * 2. / 3.
        elif self.m_kernel == 'x2':
            return pde - 4. * self.m_lamb * points / 3.
        return pde

    @classmethod
    def x2_exact(cls, points):
        """ suite of exact solutions. """

        return points ** 2

    def x2_load(self, points, time=0):
        """ Complement the pide to test algorithm convergence, for u = x**2.
        """

        pde = 2. * (self.m_diff + self.m_conv * points) + (self.m_reac
                - self.m_lamb) * points ** 2

        if self.m_kernel == '1':
            return pde + self.m_lamb * 2. / 3.
        elif self.m_kernel == 'x':
            pde += self.m_lamb * 2. * points / 3.
        elif self.m_kernel == 'x2':
            pde += self.m_lamb * (2. / 3. * points ** 2 + 2. / 5.)
        return pde

    def setUp(self):
        """ Initialize all scalar paramters.
            Note: All the convolutions were computed with respect to
                  the interval [-1, 1]. If you want to change the test cases,
                  please recompute the loads.
        """

        self.m_sigma = 0.2
        self.m_mu = 0.05

        # grid and number of segments.

        (self.m_nsegs, self.m_xl, self.m_xr) = (20, -1., 1.)

        # PIDE coefficients.

        (self.m_diff, self.m_conv, self.m_reac, self.m_lamb) = (-.30,
                1., 1., .50)

        self.m_kernel = 'x2'

        # shape functions and quadratures

        self.m_quad = LegendreGaussQuadrature(3)
        self.m_shape_funct_p2 = LagrangeP2()
        self.m_shape_funct_p1 = LagrangeP1()

        self.m_error_tol = 1.0e-03

        self.m_shape_funct = self.m_shape_funct_p2

    @staticmethod
    def _comp_errors(grid, sfuncts, numeric_sol, exact, time=0,
            plot_results=False):
        """ reconstruct to plot """

        recons = FemToReal(grid, sfuncts)
        solu = recons.reconstruct(numeric_sol)
        if time == 0:
            y_vals = exact(solu[:, 0])
        else:
            y_vals = exact(solu[:, 0], time)

        if plot_results:
            axis = plt.figure(figsize=(11.5, 5.)).gca()
            axis.plot(solu[:, 0], solu[:, 1], '-')
            plt.show()
        return norm(solu[:, 1].flatten() - y_vals.flatten())


class StaticPIDE(PIDETestSuite):

    """ Make sure matrices are assemble correctly. """

    def runTest(self):
        """ compute solution to the bvp and compare exact solution. """

        # instantiate an object to type PIDEStatic1D

        pide_engine = PIDEStatic1D(self.m_nsegs, self.m_xl, self.m_xr)

        # set reference finte element space and quad formulas.

        pide_engine.set_sfunc_and_quads(self.m_shape_funct, self.m_quad)

        # set paramters, load and integral kernel.

        self.m_diff, self.m_conv, self.m_reac = 0, 0, 0
        pide_engine.set_pide_cos(self.m_diff, self.m_conv, self.m_reac,
                                 self.m_lamb)

        # Choose the test.

        load = self.x2_load
        exact = self.x2_exact

        # set exact solution and loads.

        pide_engine.set_load(load)
        pide_engine.set_kernel(self.exact_kernel)

        # set boundary condtitions.

        pide_engine.set_bd_conditions(exact(self.m_xl),
                exact(self.m_xr))

        # solve the pide, returns a vector in fem space.

        numeric_sol = pide_engine.solve()

        # compute the error.

        grid = pide_engine.get_mesh()
        error_norm = self._comp_errors(grid, self.m_shape_funct,
                numeric_sol, exact)

        self.assert_(1.0e-02 > abs(error_norm))


class DiffusionJump1D(PIDETestSuite):

    """ The most important engine in the Levy pricing suite. """

    def exact_sol(self, x, time=0):
        """ exact solutions """

        if 0 == self.test_id:
            return exp((self.m_diff + self.m_reac) * time) * exp(-x)
        elif 1 == self.test_id:
            return time * ones(x.shape)
        elif 2 == self.test_id:
            return time * x
        elif 3 == self.test_id:
            return time * x ** 2
        elif 4 == self.test_id:
            return exp(self.m_diff * time) * exp(-x)

        assert False

    def load_k1(self, x, curr_time, next_time):
        """ load for \partial_t + a*\partial_{xx} + \lambda * {k * u - u}.
            k = 1.
        """

        if 0 == self.test_id:
            return zeros(x.shape)

        elif 1 == self.test_id:
            tmp = next_time * (self.m_reac - self.m_lamb)
            tmp += curr_time * self.m_lamb * (self.m_xr - self.m_xl)
            return 1. - tmp * ones(x.shape)

        elif 2 == self.test_id:
            tmp = next_time * (self.m_reac - self.m_lamb) * x
            tmp += next_time * self.m_conv
            tmp += curr_time / 2. * self.m_lamb * ((self.m_xr - x) ** 2
                    - (self.m_xl - x) ** 2)
            return x - tmp

        elif 3 == self.test_id:
            tmp = next_time * (self.m_reac - self.m_lamb) * x ** 2
            tmp += next_time * 2. * (self.m_diff + self.m_conv * x)
            tmp += curr_time / 3. * self.m_lamb * ((self.m_xr - x) ** 3
                    - (self.m_xl - x) ** 3)
            return x ** 2 - tmp

        elif 4 == self.test_id:
            tmp2 = (self.m_diff * next_time ** 2 - (x + self.m_lamb)) \
                * exp(-next_time * x)
            if 0 == curr_time:
                tmp2 -= self.m_lamb * (self.m_xr - self.m_xl)
            else:
                tmp2 -= self.m_lamb * exp(curr_time * x) * (exp(-curr_time
                        * self.m_xr) - exp(-time * self.m_xl)) / curr_time
            return tmp2
        assert False

    def runTest(self):
        """ compute solution to the initial value problem. """

        print ' '
        ntime_steps = 2
        time2simulate = .5

        self.test_id = 2
        self.m_kernel = '1'

        self.m_nsegs = 2
        pide_solver = PIDESolver(self.m_nsegs, self.m_xl, self.m_xr)

        pide_solver.set_sfunc_and_quads(self.m_shape_funct, self.m_quad)
        self.m_diff = 0.
        self.m_conv = 0.
        self.m_reac = 0.
        self.m_lamb = 0.5 #1.5
        pide_solver.set_pide_cos(self.m_diff, self.m_conv, self.m_reac,
                                 self.m_lamb)

        pide_solver.set_ivalf(self.exact_sol)
        pide_solver.set_lbc(self.exact_sol)
        pide_solver.set_rbc(self.exact_sol)
        pide_solver.set_kernel(self.exact_kernel)
        pide_solver.set_load(self.load_k1)

        time_stepper = IMEXTimeStepper()
        time_stepper.set_dolog(True)
        time_stepper.set_method('euler')
        time_stepper.set_number_steps(ntime_steps)

        time_stepper.set_startat(0.5)
        time_stepper.set_time2simulate(time2simulate)

        pide_solver.set_time_stepper(time_stepper)

        numeric_sol = pide_solver.solve()
        grid = pide_solver.get_mesh()

        error_norm = self._comp_errors(
            grid,
            self.m_shape_funct,
            numeric_sol,
            self.exact_sol,
            time2simulate,
            True,
            )

        self.assert_(1.0e-02 > abs(error_norm))


def my_suite():
    """ add test here. """

    suite = unittest.TestSuite()
    #suite.addTest(StaticPIDE())
    suite.addTest(DiffusionJump1D())

    return suite


if __name__ == '__main__':

    Runner = unittest.TextTestRunner(verbosity=2)
    result = Runner.run(my_suite())

